import uuid
import json
from redis.asyncio import Redis
from config.env import RedisConfig
from ocr_core.schema import OcrTask, OcrTaskResult, TaskStatus
from utils.log_util import logger
from config.mongodb import OcrTaskMDB, OcrResultMDB


class OcrProducer:
    """OCR 任务生产者，负责推送任务到 Redis 队列并查询结果"""

    def __init__(self, redis_client: Redis):
        self.redis: Redis = redis_client

    async def create_ocr_task(self, task: OcrTask) -> str:
        """
        创建新 OCR 任务并推送到 Redis 队列
        Args:
            task: OcrTask
        Returns:
            task_id: 任务ID
        """
        # 队列检查
        await self._check_queue_capacity()

        # 生成任务 ID
        task.task_id = task.task_id or str(uuid.uuid4())

        try:
            serialized_task = json.dumps(task.model_dump(), ensure_ascii=False)
        except Exception as e:
            logger.error(f"OCR 任务序列化失败: {task}, 错误: {e}")
            raise

        # 初始化任务状态
        init_result = OcrTaskResult(
            task_id=task.task_id,
            scenario=task.scenario,
            texts=[],
            img_ms=0,
            ocr_ms=0,
            llm_ms=0,
            status=TaskStatus.PENDING,
            error="",
        )

        success = await self.redis.set(
            f"{RedisConfig.redis_ocr_result}{task.task_id}",
            json.dumps(init_result.model_dump(), ensure_ascii=False),
            ex=3600,
        )

        if not success:
            logger.error(f"OCR 任务 {task.task_id} 状态初始化失败")
            raise RuntimeError("任务状态初始化失败")

        # 推送任务到队列
        await self.redis.rpush(RedisConfig.redis_ocr_queue, serialized_task)
        logger.debug(f"OCR 任务 {task.task_id} 创建成功 ✅")

        # 插入mongodb的ocr文档
        OcrTaskMDB.insert_one(task.model_dump())
        OcrResultMDB.insert_one(init_result.model_dump())
        return task.task_id

    async def get_ocr_task_result(self, task_id: str) -> OcrTaskResult:
        """
        获取任务状态
        Args:
            task_id: 任务ID
        Returns:
            OcrTaskResult
        """
        find_one = OcrResultMDB.find_one({"task_id": task_id})
        if find_one:
            try:
                return OcrTaskResult(**find_one)
            except Exception as e:
                logger.error(f"OCR 任务结果反序列化失败 [task_id={task_id}] 错误: {e}")

        # 兜底返回
        return OcrTaskResult(
            task_id=task_id,
            scenario="",
            texts=[],
            img_ms=0,
            ocr_ms=0,
            llm_ms=0,
            status=TaskStatus.NOT_FOUND,
            error="Task not found",
        )

    async def _check_queue_capacity(self) -> None:
        """检查 OCR 和 LLM 队列容量"""
        pipeline = self.redis.pipeline(False)
        pipeline.llen(RedisConfig.redis_ocr_queue)
        pipeline.llen(RedisConfig.redis_llm_queue)
        ocr_length, llm_length = await pipeline.execute()

        if ocr_length >= RedisConfig.ocr_max_queue_length:
            raise RuntimeError("OCR 任务队列已满，请稍后再试")
        if llm_length >= RedisConfig.llm_max_queue_length:
            raise RuntimeError("LLM 任务队列已满，请稍后再试")
